package com.archgeek.calcite;

import org.apache.calcite.adapter.csv.CsvSchema;
import org.apache.calcite.adapter.csv.CsvTable;
import org.apache.calcite.plan.hep.HepPlanner;
import org.apache.calcite.plan.hep.HepProgram;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.JoinRelType;
import org.apache.calcite.rel.externalize.RelWriterImpl;
import org.apache.calcite.rel.rules.FilterJoinRule;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.sql.SqlExplainLevel;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.tools.FrameworkConfig;
import org.apache.calcite.tools.Frameworks;
import org.apache.calcite.tools.RelBuilder;

import java.io.File;
import java.io.PrintWriter;

/**
 * @author pizhihui
 * @date 2024-05-30 19:02
 */
public class C3_RBO {

    public void run() throws Exception{

        /**
         * ================ 构建 schema
         */
        SchemaPlus schema = Frameworks.createRootSchema(true);
        File file = new File(getClass().getClassLoader().getResource("").getPath());

        CsvSchema csvSchema = new CsvSchema(file, CsvTable.Flavor.SCANNABLE);
        schema.add("orders", csvSchema.getTable("orders"));
        schema.add("consumers", csvSchema.getTable("consumers"));

        System.out.println(schema.getTableNames());


        FrameworkConfig config = Frameworks.newConfigBuilder().defaultSchema(schema).build();


        RelBuilder relBuilder = RelBuilder.create(config);

        RelNode relNode = relBuilder
                .scan("consumers")
                .scan("orders")
                .join(JoinRelType.INNER,
                        relBuilder.call(SqlStdOperatorTable.EQUALS,
                                relBuilder.field("id"),
                                relBuilder.field("user_id"))
                       )
                .filter(
                        relBuilder.call(SqlStdOperatorTable.EQUALS,
                                relBuilder.field("lastname"),
                                relBuilder.literal("jacky")))
                .project(
                        relBuilder.field("id"),
                        relBuilder.field("goods"),
                        relBuilder.field("price"),
                        relBuilder.field("firstname"),
                        relBuilder.field("lastname"))
                .sortLimit(0, 5, relBuilder.field("id"))
                .build();


        // System.out.println(RelOptUtil.toString(relNode));

        relNode.explain(new RelWriterImpl(new PrintWriter(System.out, true), SqlExplainLevel.ALL_ATTRIBUTES, true));


        System.out.println("============== RBO 优化 ==================");

        HepProgram hepProgram = HepProgram.builder()

        .addRuleInstance(
                FilterJoinRule.FilterIntoJoinRule.FilterIntoJoinRuleConfig.DEFAULT.toRule())
                .build();

        HepPlanner hepPlanner = new HepPlanner(hepProgram);
        hepPlanner.setRoot(relNode);
        RelNode r = hepPlanner.findBestExp();


//        System.out.println(RelOptUtil.toString(r));
//        System.out.println("=================");
//        System.out.println(r.explain());
//        System.out.println("=================");

        r.explain(new RelWriterImpl(new PrintWriter(System.out, true), SqlExplainLevel.ALL_ATTRIBUTES, true));
    }


    public static void main(String[] args) throws Exception {
        C3_RBO runner = new C3_RBO();
        runner.run();
    }


}
